import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
import os
import requests
import os.path

# 获取当前文件的目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

def download_and_process_data():
    # 创建数据目录
    if not os.path.exists(os.path.join(BASE_DIR, 'data')):
        os.makedirs(os.path.join(BASE_DIR, 'data'))

    # 下载数据
    url = "https://archive.ics.uci.edu/ml/machine-learning-databases/heart-disease/processed.cleveland.data"

    # 定义列名
    columns = [
        'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs', 'restecg',
        'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal', 'target'
    ]

    try:
        # 使用requests下载数据
        response = requests.get(url)
        response.raise_for_status()  # 检查下载是否成功

        # 将数据保存为临时文件
        with open(os.path.join(BASE_DIR, 'data', 'temp_data.txt'), 'w') as f:
            f.write(response.text)

        # 读取数据
        df = pd.read_csv(os.path.join(BASE_DIR, 'data', 'temp_data.txt'), names=columns, na_values='?')

        # 删除临时文件
        os.remove(os.path.join(BASE_DIR, 'data', 'temp_data.txt'))

        # 处理缺失值
        df = df.dropna()

        # 保存完整数据集
        df.to_csv(os.path.join(BASE_DIR, 'data', 'heart_disease_full.csv'), index=False)

        # 分割训练集和测试集
        train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

        # 保存训练集和测试集
        train_df.to_csv(os.path.join(BASE_DIR, 'data', 'heart_disease_train.csv'), index=False)
        test_df.to_csv(os.path.join(BASE_DIR, 'data', 'heart_disease_test.csv'), index=False)

        print("数据集下载和处理完成！")
        print(f"总样本数: {len(df)}")
        print(f"训练集样本数: {len(train_df)}")
        print(f"测试集样本数: {len(test_df)}")
        print("\n数据预览:")
        print(df.head())

    except Exception as e:
        print(f"下载数据时出错: {str(e)}")
    df = pd.read_csv(os.path.join(BASE_DIR, 'data', 'heart_disease_full.csv'))
    print("各类别数量：")
    print(df['target'].value_counts())
    print("\n各类别比例：")
    print(df['target'].value_counts(normalize=True))

if __name__ == "__main__":
    download_and_process_data() 